
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import os 
import itertools
import tqdm
import time
import warnings
import json
warnings.filterwarnings('ignore')

class Feature_display(): 
    close = pd.read_pickle(r'data\stock_data\daily\consentrate_price\close.pkl.gzip').astype(np.float32)
    lagret = pd.read_pickle(r'data\stock_data\daily\consentrate_price\lagret.pkl.gzip').astype(np.float32)
    
    def __init__(self, root, feature_name, feature_root='data/cmodty/my_feature/feature_data'):
        self.root = f'{root}/{feature_name[:-9]}'
        self.feature_name =feature_name
        self.feature_root = feature_root
        if not os.path.isdir(self.root):
            os.makedirs(self.root)
        self.close = Feature_display.close
        self.lagret = Feature_display.lagret
                     
    def get_data(self):
        feature = pd.read_pickle(os.path.join(self.feature_root, self.feature_name)).astype(np.float32)
        return feature 

    def data_standard(self,data):
        data_value = data.values
        data_value_std = -1 + (2/(np.nanmax(data_value, axis=0) - np.nanmin(data_value, axis=0))) * (data_value  - np.nanmin(data_value, axis=0))
        std_df = pd.DataFrame(data_value_std)
        std_df.index = data.index
        std_df.columns = data.columns 
        return std_df
   
    def describe_signal(self, data):
        perc =[0.01, 0.05 , 0.25, 0.5, 0.75, 0.95, 0.99 ]
        desc = data.describe(percentiles = perc)
        desc.to_csv(self.root+ '/' +self.feature_name+'_describe.csv')     
        desc = pd.DataFrame(desc.values.T, index=desc.columns, columns=desc.index)
        desc_desc = desc.describe(percentiles = [0.01, 0.05,  0.5,  0.95, 0.99 ])
        desc_desc.to_csv(self.root+ '/' + self.feature_name+'_describe_describe.csv')
        print('线图绘制完成')
   
    
    def nan_check(self, lagret, signal):#df
        nan_lagret = np.isnan(lagret)
        nan_signal = np.isnan(signal)
        ratio_nan_signal = np.nansum(nan_signal, axis=0)/ len(lagret)
        ratio_nan_lagret = np.nansum(nan_lagret, axis=0)/ len(lagret)
        nan_ratio = pd.DataFrame(ratio_nan_signal-ratio_nan_lagret, index=lagret.columns)
        nan_ratio.to_csv(self.root+ '/' + self.feature_name+'_nan_ratio.csv')
        print('nan值检查完成')
    
        
    def inf_check(self, lagret, signal):#df
        inf_lagret = np.isinf(lagret)
        inf_signal = np.isinf(signal)
        ratio_inf_signal = np.nansum(inf_signal, axis=0)/ len(lagret)
        ratio_inf_lagret = np.nansum(inf_lagret, axis=0)/ len(lagret)
        inf_ratio = pd.DataFrame(ratio_inf_signal-ratio_inf_lagret, index=lagret.columns)
        inf_ratio.to_csv(self.root+ '/' + self.feature_name+'_inf_ratio.csv')
        print('inf值检查完成')
    
    
    def span_std_check(self, signal, close):
        signal_std = []
        close_std = []
        index_list = []
        length = len(signal)
        span_time = length // 10
        for i in range(9):
            signal_i = signal.iloc[span_time*i: span_time*(i+1)]
            close_i = close.iloc[span_time*i: span_time*(i+1)]
            signal_std.append(signal_i.std(axis=0))
            close_std.append(close_i.std(axis=0))
            index_list.append(signal.index[span_time*i])
        close_std_df = pd.concat(close_std, axis=1)
        signal_std_df = pd.concat(signal_std, axis=1)
        close_std_df.columns = index_list
        signal_std_df.columns = index_list
        close_std_df.to_csv(self.root+ '/' + self.feature_name+'_close_spanstd.csv')
        signal_std_df.to_csv(self.root+ '/' + self.feature_name+'_signal_spanstd.csv')
    
                     
    def plot_line(self, signal, close, lagret, start_day='2021-03-01',end_day='2022-08-04',col=['600151'], std_command = False):
        signal = signal[(signal.index>=start_day)&(signal.index<end_day)]
        close = close[(close.index>=start_day)&(close.index<end_day)]
        lagret = lagret[(lagret.index>=start_day)&(lagret.index<end_day)]
        if std_command == True:
            signal = self.data_standard(signal)
            close = self.data_standard(close)
            lagret = self.data_standard(lagret)
        plt.figsize = (60, 30)
        fig, ax =  plt.subplots(3,1)            
        tick_spacing = 225        #通过修改tick_spacing的值可以修改x轴的密度
        ax[0].plot(signal[col].values, label='signal')
        ax[0].xaxis.set_major_locator(ticker.MultipleLocator(tick_spacing ))
        ax[0].legend()
        ax[1].plot(close[col].values, label='close') 
        ax[1].xaxis.set_major_locator(ticker.MultipleLocator(tick_spacing ))
        ax[1].legend()
        ax[2].plot(lagret[col].values, label='lagret') 
        ax[2].xaxis.set_tick_params(rotation=35)
        ax[2].legend() 
        plt.tight_layout()
        plt.savefig(self.root+ '/' + start_day +'-'+ end_day+ 'line.png')
        plt.close()

    
    def plot_compare(self, signal, close, lagret, start_day='2021-03-01',end_day='2022-08-04',col=['600151'], std_command = True):
        signal = signal[(signal.index>=start_day)&(signal.index<end_day)]
        close = close[(close.index>=start_day)&(close.index<end_day)]
        close[col].to_csv(self.root+ '/' + start_day +'-'+ end_day+'close.csv')
        lagret = lagret[(lagret.index>=start_day)&(lagret.index<end_day)]
        if std_command == True:
            signal = self.data_standard(signal)
            close = self.data_standard(close)
            lagret = self.data_standard(lagret)
        plt.figsize = (60, 30)
        fig, ax =  plt.subplots(2,1)            
        tick_spacing = 225       #通过修改tick_spacing的值可以修改x轴的密度
        ax[0].plot(signal[col].values, label='signal')
        ax[0].plot(close[col].values, label='close') 
        ax[0].xaxis.set_major_locator(ticker.MultipleLocator(tick_spacing ))
        ax[0].legend()

        ax[1].plot(signal[col].values, label='signal')
        ax[1].plot(lagret[col].values, label='lagret') 
        ax[1].xaxis.set_tick_params(rotation=35)
        ax[1].legend() 
        plt.tight_layout()
        plt.savefig(self.root+ '/' + start_day +'-'+ end_day+ 'conpare_line.png')
        plt.close()

    
    def plot_hist(self, signal, close, lagret,col=['600151'], std_command = False, time_inv = 1):
        plt.figsize = (20, 30)
        fig, ax =  plt.subplots(2,1)

        ax[0].hist(signal[col].values, bins =200,alpha=0.5,color = 'gold', label='signal')       
        ax[0].legend()
        ax[1].hist(lagret[col].values, bins =200,alpha=0.5,color = 'blue', label='lagret') 
        ax[1].legend()
        plt.savefig(self.root+ '/' +self.feature_name+ 'hist_.png')
        plt.close()

        plt.figsize = (20, 20)
        fig, ax = plt.subplots(nrows=1, ncols=2)
        bp1= ax[0].boxplot(signal[col].values,labels = ['15'],sym='r')
        ax[0].set_title(col[0], fontsize=20)  # 标题，并设定字号大小
        bp2= ax[1].boxplot(signal['KEE_Comdty'].values,labels = ['15'],sym='r')
        ax[1].set_title('KEE_Comdty', fontsize=20) 
        plt.savefig(self.root+ '/' +self.feature_name+ 'bins.png')
        plt.close()

    def test(self,start_day='2021-03-01',end_day='2022-08-04'):
        signal = self.get_data() #
        self.span_std_check(signal, self.close)
        self.describe_signal(signal)
        self.nan_check(self.lagret, signal)
        self.inf_check(self.lagret, signal)
        self.plot_line(signal, self.close, self.lagret, start_day, end_day)
        self.plot_compare(signal, self.close, self.lagret, start_day, end_day)  
        try:  
            self.plot_hist(signal, self.close, self.lagret)  
        except:
            print('plot hist failed')   


class IC_display(): 
    lagret = pd.read_pickle(r'data\stock_data\daily\consentrate_price\lagret.pkl.gzip').astype(np.float32)
    lagret1 = lagret.shift(-1)
    lagret10 = lagret.rolling(5).sum().shift(-5)
    lagret30 = lagret.rolling(10).sum().shift(-10)

    def __init__(self, start_date='2018-01-01', end_date='2023-01-01', feature_root=None, method='pearson') -> None:
        if feature_root is None:
            feature_root = [r'data\cmodty\my_feature\feature_data_fd']
        self.info_colums = ['nan_ratio', 'n_l', 'ic1', 'ic10', 'ic30', 'std', 'mm']
        self.save_root = r'data\cmodty\check'
        if not os.path.exists(self.save_root):
            os.makedirs(self.save_root)
        self.start_date = start_date
        self.end_date = end_date
        self.feature_root = feature_root
        self.method = method
        self.lagret = IC_display.lagret[(IC_display.lagret.index>=self.start_date)&(IC_display.lagret.index<=self.end_date)]
        self.lagret1 = IC_display.lagret1[(IC_display.lagret1.index>=self.start_date)&(IC_display.lagret1.index<=self.end_date)]
        self.lagret10 = IC_display.lagret10[(IC_display.lagret10.index>=self.start_date)&(IC_display.lagret10.index<=self.end_date)]
        self.lagret30 = IC_display.lagret30[(IC_display.lagret30.index>=self.start_date)&(IC_display.lagret30.index<=self.end_date)]
    
    def record_ic(self, lagret, signal):
        lagret = lagret[(lagret.index>=self.start_date)&(lagret.index<=self.end_date)]
        signal = signal[(signal.index>=self.start_date)&(signal.index<=self.end_date)]
        ic = self.ic_calc_func(lagret.values, signal.values)
        return ic

    def ic_calc_func(self,array_a, array_b):
        array_a[np.isinf(array_a)] = np.nan
        array_b[np.isinf(array_b)] = np.nan
        if self.method == 'spearman':
            array_a = array_a.reshape((-1,))
            array_b = array_b.reshape((-1,))
            value = np.stack([array_a,array_b],axis=1)
            df = pd.DataFrame(value)
            ic_mat = df.rank().corr(min_periods=1)
        else:
            array_a = array_a.reshape((-1,), order='F')
            array_b = array_b.reshape((-1,), order='F')
            value = np.stack([array_a,array_b],axis=1)
            df = pd.DataFrame(value)
            ic_mat = df.corr(min_periods=1)
        ic = ic_mat.values[0,1]
        return ic

    def record_nan_ratio(self, lagret, signal):
        lagret_nan = self.nan_calc_func(lagret.values)
        signal_nan = self.nan_calc_func(signal.values)
        return signal_nan, signal_nan-lagret_nan

    def nan_calc_func(self, fv):
        fv = fv.reshape((-1,), order='F')
        fv[np.isinf(fv)] = np.nan
        nan_ratio = np.isnan(fv).sum()/ len(fv)
        return nan_ratio

    def recore_mm(self, signal):
        return self.mm_calc_func(signal.values)
        
    def mm_calc_func(self, fv):
        fv[np.isinf(fv)] = np.nan
        fv = fv.reshape((-1,), order='F')
        mm = np.nanmax(fv)
        std = np.nanstd(fv)
        return mm, std

    def info_generate(self):
        index_lst = []
        nan_ratio_lst = []
        n_l_lst = []
        ic1_lst = []
        ic10_lst = []
        ic30_lst = []
        mm_lst = []
        std_lst = []
        for j in self.feature_root:
            lst = sorted(os.listdir(j))
            for i in tqdm.tqdm(lst):
                try:
                    name_i = i[:-9]
                    path = os.path.join(j, i)
                    feature_i = pd.read_pickle(path)
                    feature_i = feature_i[(feature_i.index>=self.start_date)&(feature_i.index<=self.end_date)]
                    feature_i = feature_i.astype(np.float32)
                    nan_ratio, n_l = self.record_nan_ratio(self.lagret, feature_i)
                    ic1 = self.record_ic(self.lagret1, feature_i)
                    ic10 = self.record_ic(self.lagret10, feature_i)
                    ic30 = self.record_ic(self.lagret30, feature_i)
                    mm, std = self.recore_mm(feature_i)
                    nan_ratio_lst.append(nan_ratio) 
                    n_l_lst.append(n_l)
                    ic1_lst.append(ic1)
                    ic10_lst.append(ic10)
                    ic30_lst.append(ic30)
                    mm_lst.append(mm)
                    std_lst.append(std)
                    index_lst.append(name_i)
                    # print(f'{name_i}记录完成')
                except:
                    print(f'error {i}')
                    continue
        value = np.stack([nan_ratio_lst, n_l_lst, ic1_lst, ic10_lst, ic30_lst, mm_lst, std_lst], axis=1)
        df = pd.DataFrame(value, columns=self.info_colums, index=index_lst)
        tail = self.get_tail()
        df.to_csv(os.path.join(self.save_root, f'{self.method}_info_{tail}.csv'))

    def get_tail(self):
        tail_lst = []
        for i in self.feature_root:
            tail_i = os.path.split(i)[-1]
            tail_lst.append((tail_i.split('_')[-1]))
        tail = ''
        for i in tail_lst:
            tail = tail+i
        return tail
        

class Cor_display_fast():
    def __init__(self, start_date='2020-01-01', end_date='2023-01-01', feature_root=None, thred=0.85, cor_thred =0.01, nan_thred =0.05, method='pearson') -> None:
        if feature_root is None:
            feature_root = [r'data\cmodty\my_feature\feature_data_fd']
        self.save_root = r'data\cmodty\check'
        if not os.path.exists(self.save_root):
            os.makedirs(self.save_root)
        self.start_date = start_date
        self.end_date = end_date
        self.feature_root = feature_root
        self.thred = thred
        self.cor_thred = cor_thred
        self.nan_thred = nan_thred
        self.tail = self.get_tail()
        self.method = method

    def value_reshape(self, signal):
        signal = signal[(signal.index>=self.start_date)&(signal.index<=self.end_date)]
        array_b = signal.values
        array_b[np.isinf(array_b)] = np.nan
        array_b = array_b.reshape((-1,), order='F')
        return array_b

    def cor_reshape(self, df):
        name1_lst = []
        name2_lst = []
        cor_lst = []
        col = df.columns
        idx = df.index
        for i in range(len(col)):
            for j in range(len(idx)):
                if i < j:
                    name1_lst.append(col[i])
                    name2_lst.append(idx[j])
                    cor_lst.append(df.iloc[i,j])
        value = np.stack([name1_lst, name2_lst, cor_lst], axis=1)
        df = pd.DataFrame(value, columns=['name1','name2','cor'])
        df.to_csv(os.path.join(self.save_root,f'cor_raw_{self.tail}.csv')) 

    def feature_select(self, matrix):
        org_list = (matrix.columns).tolist()
        column_list = org_list.copy()
        num = len(column_list)
        drop_list = [np.nan]
        matrix_value = matrix.values
        for i, j in itertools.product(range(num), range(num)):
            if (  i < j and (np.abs(matrix_value[i, j]) > self.thred)):
                if (column_list[i] not in drop_list) :
                    drop_list.append(column_list[j])

        column_list = list(set(org_list)-set(drop_list))
        feature_fliter_result = {
            'origin': org_list,
            'drop': drop_list,
            'reamain': column_list,
        }
        feature_result = json.dumps(feature_fliter_result,indent=4)
        tail = self.get_tail()
        with open(os.path.join(self.save_root,f'fliter_{tail}.json') , 'w') as f:
            f.write(feature_result)
        # df = pd.DataFrame(np.stack([org_list, drop_list, column_list], axis=1), columns=['orgin', 'drop', 'remain'])
        df = pd.concat([pd.DataFrame(org_list), pd.DataFrame(drop_list), pd.DataFrame(column_list)], axis=1, ignore_index=True)
        df.columns = ['orgin', 'drop', 'remain']
        df.to_csv(os.path.join(self.save_root,f'fliter_{tail}.csv'))

        return column_list               

    def cor_generate(self):
        name_lst = []
        tail = self.get_tail()
        info_df = self.get_info_df()
        info_select = info_df[((abs(info_df['ic1'])>self.cor_thred)|(abs(info_df['ic10'])>self.cor_thred)|(abs(info_df['ic30'])>self.cor_thred))&((info_df['nan_ratio']<self.nan_thred)&(info_df['std']>0))]
        info_select.to_csv(os.path.join(self.save_root,f'info_select_{tail}.csv'))
        info_id = info_select.index
        df_lst = []
        for j in self.feature_root:
            lst = sorted(os.listdir(j))
            for i in tqdm.tqdm(lst):
                if i[:-9] in info_id:
                    name_lst.append(i[:-9])
                    path = os.path.join(j, i)
                    # file_lst.append(path)
                    featture_i = pd.read_pickle(path).astype(np.float16)
                    fv = self.value_reshape(featture_i)
                    df_lst.append(fv)
        value = np.stack(df_lst, axis=1)
        df = pd.DataFrame(value, columns=name_lst)
        cor = df.corr(min_periods=1)
        cor.to_csv(os.path.join(self.save_root,f'cor_{self.tail}.csv'))
        self.feature_select(cor)
        self.cor_reshape(cor)
        # 
    def get_tail(self):
        tail_lst = []
        for i in self.feature_root:
            tail_i = os.path.split(i)[-1]
            tail_lst.append((tail_i.split('_')[-1]))
        tail = ''
        for i in tail_lst:
            tail = tail+i
        return tail

    def get_info_df(self):
        info_lst = []
        for i in self.feature_root:
            tail_i = os.path.split(i)[-1]
            tail = (tail_i.split('_'))[-1]
            info_i = pd.read_csv(os.path.join(self.save_root,f'{self.method}_info_{tail}.csv'), index_col=0)
            info_lst.append(info_i)
        df = pd.concat(info_lst, axis=1)
        return df


class Cor_display_straitfy(Cor_display_fast):
    def __init__(self, start_date='2020-01-01', end_date='2023-01-01', feature_root=None, thred=0.85, cor_thred=0.01, nan_thred=0.05, method='pearson') -> None:
        super().__init__(start_date, end_date, feature_root, thred, cor_thred, nan_thred, method)
    

    def cor_generate(self):
        name_lst = []
        info_df = self.get_info_df()
        tail =self.get_tail()
        info_select = info_df[((abs(info_df['ic1'])>self.cor_thred)|(abs(info_df['ic10'])>self.cor_thred)|(abs(info_df['ic30'])>self.cor_thred))&((info_df['nan_ratio']<self.nan_thred)&(info_df['std']>0))]
        info_select.to_csv(os.path.join(self.save_root,f'info_select_{tail}.csv'))
        info_id = info_select.index
        df_lst = []
        remain_list = self.get_first_remain()
        for j in self.feature_root:
            lst = sorted(os.listdir(j))
            for i in tqdm.tqdm(lst):
                if (i[:-9] in info_id) and (i[:-9] in remain_list):
                    name_lst.append(i[:-9])
                    path = os.path.join(j, i)
                    # file_lst.append(path)
                    featture_i = pd.read_pickle(path).astype(np.float16)
                    fv = self.value_reshape(featture_i)
                    df_lst.append(fv)
        value = np.stack(df_lst, axis=1)
        df = pd.DataFrame(value, columns=name_lst)
        cor = df.corr(min_periods=1)
        cor.to_csv(os.path.join(self.save_root,'cor.csv'))
        self.feature_select(cor)
        self.cor_reshape(cor)

    def get_first_remain(self):
        remain_lst = []
        for i in self.feature_root:
            tail_i = os.path.split(i)[-1]
            tail= (tail_i.split('_'))[-1]
            with open(os.path.join(self.save_root,f'fliter_{tail}.json'),'r', encoding='UTF-8') as f:
                self.select_dict = json.load(f)
                select_lst = self.select_dict['reamain']
                remain_lst.extend(select_lst)
        return remain_lst

    def get_info_df(self):
        info_lst = []
        for i in self.feature_root:
            tail_i = os.path.split(i)[-1]
            tail = (tail_i.split('_'))[-1]
            info_i = pd.read_csv(os.path.join(self.save_root,f'{self.method}_info_{tail}.csv'), index_col=0)
            info_lst.append(info_i)
        df = pd.concat(info_lst, axis=0)
        return df
        


    


# class Cor_display():
#     def __init__(self, start_date='2020-01-01', end_date='2023-01-01', feature_root=None) -> None:
#         if feature_root is None:
#             feature_root = [r'data\cmodty\my_feature\feature_data_fd']
#         self.save_root = r'data\cmodty\check'
#         if not os.path.exists(self.save_root):
#             os.makedirs(self.save_root)
#         self.start_date = start_date
#         self.end_date = end_date
#         self.feature_root = feature_root

#     def ic_calc_func(self,array_a, array_b):
#         array_a[np.isinf(array_a)] = np.nan
#         array_b[np.isinf(array_b)] = np.nan
#         array_a = array_a.reshape((-1,), order='F')
#         array_b = array_b.reshape((-1,), order='F')
#         ic = (np.nanmean((array_a*array_b))-np.nanmean(array_a)*np.nanmean(array_b))/(np.nanstd(array_a)*np.nanstd(array_b))
#         return ic

#     def record_ic(self, lagret, signal):
#         lagret = lagret[(lagret.index>=self.start_date)&(lagret.index<=self.end_date)]
#         signal = signal[(signal.index>=self.start_date)&(signal.index<=self.end_date)]
#         ic = self.ic_calc_func(lagret.values, signal.values)
#         return ic

#     def cor_generate(self):
#         file_lst = []
#         info_df = pd.read_csv(os.path.join(self.save_root,'pearson_info.csv'), index_col=0)
#         info_select = info_df[(abs(info_df['ic1'])>0.01)|(abs(info_df['ic10'])>0.01)|(abs(info_df['ic30'])>0.01)&(info_df['nan_ratio']<0.05)]
#         info_id = info_select.index
#         for j in self.feature_root:
#             lst = sorted(os.listdir(j))
#             for i in lst:
#                 if i[:-9] in info_id:
#                     path = os.path.join(j, i)
#                     file_lst.append(path)
#         ic_lst = []
#         name1_lst = []
#         name2_lst = []
#         for e in itertools.combinations(file_lst, 2):
#             try:
#                 t1 = time.time()
#                 data1 = pd.read_pickle(e[0]).astype(np.float32)
#                 data2 = pd.read_pickle(e[1]).astype(np.float32)
#                 t2=time.time()
#                 ic = self.record_ic(data1, data2)
#                 t3=time.time()
#                 ic_lst.append(ic)
#                 name1 = os.path.split(e[0])[-1][:-9]
#                 name2 = os.path.split(e[1])[-1][:-9]
#                 name1_lst.append(name1)
#                 name2_lst.append(name2)
#                 print(f'组合({name1},{name2})计算完成')
#             except:
#                 print(f'error {e}')
#                 continue
#         value = np.stack([name1_lst, name2_lst, ic_lst], axis=1)
#         df = pd.DataFrame(value, columns=['name1','name2','cor'])
#         df.to_csv(os.path.join(self.save_root),'cor.csv')